#include "stream_reader_ffmpeg.h"

StreamReaderFFmpeg::StreamReaderFFmpeg(std::string url, int camid)
{
  this->url = url;
  this->camid = camid;
  pthread_mutex_init(&this->lock,NULL); //init mutex
}

StreamReaderFFmpeg::~StreamReaderFFmpeg()
{
  this->release();
  pthread_mutex_destroy(&this->lock); //destroy mutex
}

int StreamReaderFFmpeg::init()
{
  std::cout << "[INFO]: Opening Stream: " << this->url << std::endl;

  char option_key[] = "rtsp_transport";
  char option_value[] = "tcp";
  av_dict_set(&this->avdic, option_key, option_value, 0);
  char option_key2[] = "stimeout";
  char option_value2[] = "5000000";
  av_dict_set(&this->avdic, option_key2, option_value2, 0);

  int ret = -1;
  if ((ret = avformat_open_input(&this->ifmt_ctx, this->url.c_str(), 0, &this->avdic)) < 0) {
    std::cerr << "[Error]: Couldn't Open Input File." << std::endl;
    this->release();
    return ret;
  }

  if ((ret = avformat_find_stream_info(this->ifmt_ctx, 0)) < 0) {
    std::cerr << "[Error]: Couldn't Find Stream Information." << std::endl;
    this->release();
    return ret;
  }
  av_dump_format(this->ifmt_ctx, 0, this->url.c_str(), 0);

  for (int i = 0; i < this->ifmt_ctx->nb_streams; i++) {
    if (this->ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
      this->video_index = i;
      std::cout << "[INFO]: Stream " << i << " Get VideoStream. " << this->url << std::endl;
      break;
    }
  }

  AVCodecParameters *codec_par = this->ifmt_ctx->streams[video_index]->codecpar;
  this->codec = avcodec_find_decoder(codec_par->codec_id);
  if (this->codec == NULL) {
    std::cerr << "[Error]: Unsupported Codec!" << std::endl;
    this->release();
    return -1;
  }
  this->codec_ctx = avcodec_alloc_context3(this->codec);
  if (this->codec_ctx == NULL) {
    std::cerr << "[Error]: Allocate AVCodecContext Failed!" << std::endl;
    this->release();
    return -1;
  }
  if (avcodec_parameters_to_context(this->codec_ctx, codec_par) < 0) {
    std::cerr << "[Error]: Fill AVCodecContext Failed!" << std::endl;
    this->release();
    return -1;
  }
  if (avcodec_open2(this->codec_ctx, this->codec, nullptr) < 0) {
    std::cerr << "[Error]: Couldn't Open Codec!" << std::endl;
    this->release();
    return -1;
  }

  this->numBytes = av_image_get_buffer_size(AV_PIX_FMT_BGR24, codec_ctx->width, codec_ctx->height, 1);
  if (this->numBytes < 0) {
    std::cerr << "[Error]: Couldn't Get the Buffer Size." << std::endl;
    this->release();
    return -1;
  }
  this->buffer = (uint8_t *)av_malloc(numBytes);
  if (this->buffer == NULL) {
    std::cerr << "[Error]: Memory Block Cannot Be Allocated!" << std::endl;
    this->release();
    return -1;
  }

  this->rawframe = av_frame_alloc();
  this->bgrframe = av_frame_alloc();
  if (this->rawframe == NULL || this->bgrframe == NULL) {
    std::cerr << "[Error]: Frame Memory Block Cannot Be Allocated!" << std::endl;
    this->release();
    return -1;
  }
  av_image_fill_arrays(this->bgrframe->data, this->bgrframe->linesize, this->buffer, AV_PIX_FMT_BGR24, this->codec_ctx->width, this->codec_ctx->height, 1);

  std::cout << "[INFO]: Open Stream Successed. " << this->url << std::endl;
  return 0;
}

void StreamReaderFFmpeg::release()
{
  if (this->img_cvrt_ctx != NULL) {
    sws_freeContext(this->img_cvrt_ctx);
    this->img_cvrt_ctx = NULL;
  }
  if (this->buffer != NULL) {
    av_free(this->buffer);
    this->buffer = NULL;
  }
  if (this->rawframe != NULL) {
    av_frame_free(&this->rawframe);
    this->rawframe = NULL;
  }
  if (this->bgrframe != NULL) {
    av_frame_free(&this->bgrframe);
    this->bgrframe = NULL;
  }
  if (this->codec_ctx != NULL) {
    avcodec_free_context(&this->codec_ctx);
    this->codec_ctx = NULL;
  }
  if (this->avdic != NULL) {
    av_dict_free(&this->avdic);
    this->avdic = NULL;
  }
  if (this->ifmt_ctx != NULL) {
    avformat_close_input(&this->ifmt_ctx);
    this->ifmt_ctx = NULL;
  }
  std::cout << "[INFO]: releasing... " << this->url << std::endl;
}

void StreamReaderFFmpeg::Play()
{
  AVPacket pkt;
  int keyval = 0, ret = 0;
  int rccount = 3;

  while (!this->stop) {
    cv::Mat image;
    ret = av_read_frame(this->ifmt_ctx, &pkt);
    if (ret < 0) {
      std::cerr << "[Error]: Read Frame Error or End Of File." << std::endl;
      //timeout or re-connect
      int connected = 0;
      for (int ri = 0; ri < rccount; ri++) {
        this->release();
        if (!(this->init()<0)) { connected = 1; break; }
      }
      if (connected) { continue; } else { break; }
    }

    if (pkt.stream_index == this->video_index) {

      pthread_mutex_lock(&this->lock); //lock
      this->curFrame = image;
      pthread_mutex_unlock(&this->lock);//unlock

      //// JUST RECV KEY FRAME
      //if (!(pkt.flags & AV_PKT_FLAG_KEY)) {
      //  continue;
      //}

      ret = avcodec_send_packet(this->codec_ctx, &pkt);
      if (ret < 0) {
        std::cerr << "[Error]: avcodec_send_packet encounter error." << std::endl;
        //timeout or re-connect
        int connected = 0;
        for (int ri = 0; ri < rccount; ri++) {
          this->release();
          if (!(this->init()<0)) { connected = 1; break; }
        }
        if (connected) { continue; } else { break; }
      }

      while (ret >= 0) {
        ret = avcodec_receive_frame(this->codec_ctx, this->rawframe);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
          break;
        }
        else if (ret < 0) {
          std::cerr << "[Error]: avcodec_receive_frame encounter error." << std::endl;
          break;
        }
	if (ret >= 0) {
          AVPixelFormat pixFormat;
          switch (this->codec_ctx->pix_fmt) {
            case AV_PIX_FMT_YUVJ420P:
              pixFormat = AV_PIX_FMT_YUV420P;
              break;
            case AV_PIX_FMT_YUVJ422P:
              pixFormat = AV_PIX_FMT_YUV422P;
              break;
            case AV_PIX_FMT_YUVJ444P:
              pixFormat = AV_PIX_FMT_YUV444P;
              break;
            case AV_PIX_FMT_YUVJ440P:
              pixFormat = AV_PIX_FMT_YUV440P;
              break;
            default:
              pixFormat = this->codec_ctx->pix_fmt;
              break;
          }
          this->img_cvrt_ctx = sws_getCachedContext(this->img_cvrt_ctx, this->codec_ctx->width, this->codec_ctx->height, pixFormat,
             this->codec_ctx->width, this->codec_ctx->height, AV_PIX_FMT_BGR24, SWS_BICUBIC, nullptr, nullptr, nullptr);
          sws_scale(this->img_cvrt_ctx, reinterpret_cast<const uint8_t *const *>(this->rawframe->data), this->rawframe->linesize, 0, this->codec_ctx->height, this->bgrframe->data, this->bgrframe->linesize);

          //cv::Mat image(this->codec_ctx->height, this->codec_ctx->width, CV_8UC3, this->bgrframe->data[0]);
          cv::Mat frame = cv::Mat(this->codec_ctx->height, this->codec_ctx->width, CV_8UC3, this->bgrframe->data[0]);

          pthread_mutex_lock(&this->lock); //lock
          frame.copyTo(this->curFrame);
          pthread_mutex_unlock(&this->lock);//unlock

          //std::cout << "Get a frame." << std::endl;
          //cv::imwrite("video.jpg", this->curFrame);
          //exit(0);
          //cv::imshow("video", this->curFrame);
          //keyval = cv::waitKey(1);
        }
      }
    }
    av_packet_unref(&pkt);
    //if (27 == keyval) break;
  }
  this->stoped = true;
}

void StreamReaderFFmpeg::Run()
{
  int ret = pthread_create(&(this->thd), NULL, StreamReaderFFmpeg::ThreadFunc, this);
  if (0 != ret) { std::cerr << "[Error]: Stream Reader Run() Failed! <source>: " << this->url << std::endl; }
  pthread_detach(this->thd);
}

void* StreamReaderFFmpeg::ThreadFunc(void* arg)
{
  StreamReaderFFmpeg *thiz = static_cast<StreamReaderFFmpeg *>(arg);
  std::cout << "[INFO]: Playing " << thiz->url << std::endl;
  thiz->Play();
  return NULL;
}

void StreamReaderFFmpeg::Stop()
{
  this->stop = true;
}

cv::Mat StreamReaderFFmpeg::GetCurFrame()
{
  cv::Mat image;
  pthread_mutex_lock(&this->lock); //lock
  this->curFrame.copyTo(image);
  pthread_mutex_unlock(&this->lock);//unlock
  return image;
}

